package org.github.manx98.jlibhac.tools.fssystem;

import lombok.Data;
import org.github.manx98.jlibhac.common.math.IByteUtil;
import org.github.manx98.jlibhac.common.math.Uint32;
import org.github.manx98.jlibhac.common.reflect.StructureMap;
import org.github.manx98.jlibhac.fs.IStorage;
import org.github.manx98.jlibhac.fs.ValueSubStorage;
import org.github.manx98.jlibhac.fssystem.BucketTree;

import java.io.IOException;

import static org.github.manx98.jlibhac.tools.fssystem.CompressionType.Lz4;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/6 14:21
 */
public class CompressedStorage extends IStorage {
    public static final int NodeSize = 0x4000;
    public static final int entrySize = IByteUtil.sizeOf(Entry.class);
    private final BucketTree bucketTree;
    private final ValueSubStorage dataStorage;

    public void initialize(ValueSubStorage dataStorage, ValueSubStorage nodeStorage, ValueSubStorage entryStorage, int bucketTreeEntryCount) throws IOException {
        nodeStorage.getSubStorage().writeAllBytes("nodeStorage");
        entryStorage.getSubStorage().writeAllBytes("entryStorage");
        bucketTree.initialize(nodeStorage, entryStorage, NodeSize, entrySize, bucketTreeEntryCount);
    }

    @Override
    public void read(long offset, byte[] destination, int destOffset, int len) throws IOException {
        // Validate arguments
        BucketTree.Offsets offsets = bucketTree.getOffsets();

        if (!offsets.isInclude(offset, destination.length))
            throw new IllegalArgumentException("OutOfRange offset: " + offset);

        // Find the offset in our tree
        var visitor = new BucketTree.Visitor();
        bucketTree.find(visitor, offset);
        long entryOffset = visitor.getEntry().VirtualOffset;
        if (entryOffset < 0 || !offsets.isInclude(entryOffset))
            throw new IllegalArgumentException("UnexpectedInCompressedStorageA: " + entryOffset);

        // Prepare to operate in chunks
        long currentOffset = offset;
        long endOffset = offset + destination.length;

        byte[] workBufferEnc = null;
        byte[] workBufferDec = null;

        while (currentOffset < endOffset) {
            // Get the current entry
            var currentEntry = visitor.getEntry();

            // Get and validate the entry's offset
            long currentEntryOffset = currentEntry.VirtualOffset;
            if (currentEntryOffset > currentOffset)
                throw new IllegalArgumentException("UnexpectedInCompressedStorageA: " + currentEntryOffset);

            // Get and validate the next entry offset
            long nextEntryOffset;
            if (visitor.canMoveNext()) {
                visitor.moveNext();
                nextEntryOffset = visitor.getEntry().VirtualOffset;
                if (!offsets.isInclude(nextEntryOffset))
                    throw new IllegalArgumentException("UnexpectedInCompressedStorageA: " + nextEntryOffset);
            } else {
                nextEntryOffset = offsets.EndOffset;
            }

            if (currentOffset >= nextEntryOffset)
                throw new IllegalArgumentException("UnexpectedInCompressedStorageA: " + nextEntryOffset);

            // Get the offset of the data we need in the entry
            long dataOffsetInEntry = currentOffset - currentEntryOffset;
            long currentEntrySize = nextEntryOffset - currentEntryOffset;

            // Determine how much is left
            long remainingSize = endOffset - currentOffset;
            long toWriteSize = Math.min(remainingSize, currentEntrySize - dataOffsetInEntry);
            if (currentEntry.CompressionType == Lz4) {
                workBufferEnc = ensureBufferSize(workBufferEnc, currentEntry.PhysicalSize.intValue());
                workBufferDec = ensureBufferSize(workBufferDec, (int) currentEntrySize);

                byte[] encBuffer = workBufferEnc;
                byte[] decBuffer = workBufferDec;

                dataStorage.read(currentEntry.PhysicalOffset, encBuffer, 0, currentEntry.PhysicalSize.intValue());

                IByteUtil.Lz4Decompress(encBuffer, decBuffer);
                System.arraycopy(decBuffer, (int) dataOffsetInEntry, destination, (int) (currentOffset - offset), (int) toWriteSize);
            } else if (currentEntry.CompressionType == CompressionType.None) {
                dataStorage.read(currentEntry.PhysicalOffset + dataOffsetInEntry, destination, (int) (currentOffset - offset), (int) toWriteSize);
            } else if (currentEntry.CompressionType == CompressionType.Zeroed) {
                for (int i = 0; i < toWriteSize; i++) {
                    destination[(int) (currentOffset - offset + i)] = 0;
                }
            }
            currentOffset += toWriteSize;
        }
    }

    static byte[] ensureBufferSize(byte[] buffer, int requiredSize) {
        if (buffer == null || buffer.length < requiredSize) {
            return new byte[requiredSize];
        }
        return buffer;
    }

    @Override
    public long getSize() {
        throw new UnsupportedOperationException("not implemented");
    }

    @Data
    public static class Entry {
        @StructureMap(offset = 0)
        public long VirtualOffset;
        @StructureMap(offset = 8)
        public long PhysicalOffset;
        @StructureMap(offset = 16)
        public CompressionType CompressionType;
        @StructureMap(offset = 17)
        public byte CompressionLevel;
        @StructureMap(offset = 18)
        public Uint32 PhysicalSize;
    }

    public CompressedStorage() {
        super(0);
        this.bucketTree = new BucketTree();
        this.dataStorage = new ValueSubStorage();
    }

    public static long queryNodeStorageSize(int entryCount) {
        return BucketTree.queryNodeStorageSize(NodeSize, entrySize, entryCount);
    }

    public static long queryEntryStorageSize(int entryCount) {
        return BucketTree.queryEntryStorageSize(NodeSize, entryCount, entryCount);
    }
}
